# src/plot_data/comparecomponent.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

"""
代码说明: 
    功能：绘制测试所有案例同一Parent的接口时间开销
    作用：逐层分析可扩展性瓶颈接口
    输入： -p parent_function(父亲接口名) -sdir store_result_dir(结果存放文件夹) -ot store_type(保存结果的类型) data_dir(数据库存放文件夹) plot_case(需要画图的算例)
    输出：store_type有,
        csv：
            csv_time
            csv_efficiency
            csv_sizeperprocess
        pic：
            pic_line
            pic_percent
            pic_efficiency
            pic_sizeperprocess
            pic_paperline
            pic_paperpercent
"""

import os
import json
import argparse
import sys 
import pandas as pd
import numpy as np

sys.path.append("..")
from package.plot_utility import get_runcases

from baseplot import plot_linechart
from baseplot import plot_percent
from baseplot import plot_sizeperprocess
from baseplot import plot_heatmap
from baseplot import casename2number
from baseplot import getplteffdata
from baseplot import mypaper_subplot
from baseplot import create_subplot
from baseplot import close_sublot
from baseplot import my_paper_plot_percent

# TODO: 这部分代码需要重构
# 1. csv和pic参数可以结合, csv_time这样形成到一起, 修改 args.plot_type 和 args.csv_type 的接口
# 2. TotalTime不是必须的, 
# 3. getpltdata完全没有必要存在

# set parser
def make_parser(): 
    parser = argparse.ArgumentParser(
        prog="jxpamg-predo-merge-pid-timerdata",
        description="jxpamg timer data preprocess script, merge pid's timer data to a single json.")
    parser.add_argument(
            "-p", "--parent",
            dest="component_parent",
            default=0,
            help="Componen under parent."
        )
    parser.add_argument(
            "-l", "--level",
            dest="out_level",
            default=-1,
            type=int,
            help="Output in level."
        )
    parser.add_argument(
            "-sdir","--savedir",
            dest="savedir",
            default="./",
            help="Destination save pic/csv directories."
        )
    parser.add_argument(
            "-ot", "--outtype",
            dest="out_type",
            default="pic_line",
            help="Set type of pic out. Option: csv_time, csv_efficiency, csv_sizeperprocess, pic_line, pic_parent, pic_efficiency, pic_sizeperprocess, pic_papreline, pic_paperpercent"
    )
    parser.add_argument(
            "LOGDIR",
            metavar="LOGDIR",
            help="Destination log directories."
        )
    parser.add_argument(
            "CASENAME",
            metavar="CASENAME",
            default="jxpamg_solver01",
            nargs='+',
            help="Set casename of plotting picture."
    )
    return parser

# 得到需要读取的文件
def getloadcase(runcases, casename):
    loadcasename = []
    for i in runcases:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        if basename in casename:
            loadcasename.append(i)
    return loadcasename

# dataprocessing
# 先处理 casename 再处理 level 和 parent, 拿到 tave
def getjsondata(casename, level, parent, plttype):
    childtimedic = {}
    totaltimedic = {}
    myplttype = plttype.split("_")[-1]
    
    for i in casename:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        with open(i, "r") as fo:
            text = fo.read()
            jsondata = json.loads(text)
            print(f"Load {basename} data.")
        if(parent):
            childtimedic[basename] = getchildavetime(jsondata, parent)
        else:
            childtimedic[basename] = getlevelavetime(jsondata, level)
        if myplttype == "total":
            totaltimedic[basename] = gettotalime(jsondata)
    return (childtimedic,totaltimedic)

# 获取相同level的时间开销
def getlevelavetime(jsondata, level):
    datalist = []
    for i in jsondata:
        tempdic = {}
        tempdic["np"] = int(i.split("_")[1])
        # paren == NULL,取得level相同的组件
        for j in jsondata[i]:
            if("level" in jsondata[i][j]):
                if(jsondata[i][j]["level"] == level ):
                    tempdic[j] = float(jsondata[i][j]["tavg"])
            else:
                print(f"{i} Name-{j} doesn't have level.")
        datalist.append(tempdic)
    # print(datalist)
    return datalist

# 获取component孩子的时间开销
def getchildavetime(jsondata, parent):
    datalist = []
    for i in jsondata:
        tempdic = {}
        tempdic["np"] = int(i.split("_")[1])
        # paren == NULL,取得level相同的组件
        for j in jsondata[i]:
            if(jsondata[i][j]["parent"] == parent):
                tempdic[j] = float(jsondata[i][j]["tavg"])
        datalist.append(tempdic)
    # print(datalist)
    return datalist

def gettotalime(jsondata):
    datalist = []
    for i in jsondata:
        tempdic = {}
        tempdic["np"] = int(i.split("_")[1])
        #paren == NULL,取得level相同的组件
        for j in jsondata[i]:
            if j == "Totaltime" :
                tempdic["TotalSolveTime"] = float(jsondata[i][j]["tavg"])
        datalist.append(tempdic)
    # print(datalist)
    return datalist

def createtitle(level, parent):
    out_title = ""
    if level >= 0:
        out_title = out_title + f"l_{level}-"
    if parent:
        out_title = out_title + f"p_{parent}-"
    return out_title

def savedata2csv(savedir, casename, out_title, type, datadict):
    for i in casename:
        casestr = casestr + casename2number(i)

    writer = pd.ExcelWriter(savedir + f"{casestr}_{out_title}_{type}.xlsx")

    k = 0
    for i in casename:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        data = pd.DataFrame(datadict[basename])
        if(type == "time"):
            data = data.sort_index(axis=1)
            data = data.sort_values(by='np', ascending=True)
        data.to_excel(writer,sheet_name=f"{casename2number(i)}")
        k += 1
    writer.save()

def save2csv(savedir,casename,out_title,csv_type,datadict):
    """表格数据处理"""
    timedict = datadict[0]
    if csv_type == "time":
        savedata2csv(savedir,casename,out_title,csv_type,timedict)
    elif csv_type == "rank_correlation":
        """
        参考文章Statistical Scalability Analysis of Communication Operations in Distributed Applications的指标
        1. 先把不同进程数下的rank算出来
        2. 再把rank与时间做线性相关
        """
        resultdic = {}
        for i in timedict:
            temp_ditc = {}
            data = pd.DataFrame(timedict[i])
            data = data.sort_index(axis=1)
            data = data.sort_values(by='np', ascending=True)
            # data = data.drop(labels="TotalSolveTime",axis=1)
            data = data.copy(deep=True)
            for col in list(data.columns):
                data[col] = data[col].rank()
            for col in list(data.columns):
                if(col != "np"):
                    
                    rs = np.cov(data["np"].values[:],data[col].values[:])[0][1] / (np.std(data["np"].values[:]) * np.std(data[col].values[:]))
                    temp_ditc[col] = rs
            resultdic[i] = [temp_ditc]
        savedata2csv(savedir,casename,out_title,csv_type,resultdic)
    return


# 画图处理的代码

# 得到画图的输入数据
# def getpltdata(data) :
#     x = []
#     for i in data["np"] :
#         x.append(i)
#     legend_name = []
#     for i in data :
#         if i != "np" :
#             #TODO:这部分代码这么写，TotalSolveTime总是被额外提出，而且应该用pandas操作
#             if i != "TotalSolveTime" :
#                 legend_name.append(i)
#             else :
#                 Total_time = list(data[i])
#     y = []
#     for i in legend_name :
#         y.append(list(data[i]))
#     return (x,y,legend_name,Total_time)

# plot
def save2pic(savedir,casename,level,parent,plttype,datadict):
    childdatadict = datadict[0]
    totaldatadict = datadict[1]

    if(plttype == "pic_paperline" ):
        fig, axes = create_subplot(nrows=2, ncols=3, xlable="Number  of  Processors", ylable="Time(s)")
    elif(plttype == "pic_paperpercent" ):
        fig, axes = create_subplot(nrows=2, ncols=3, xlable="Number  of  Processors", ylable="Percrnt(%)")

    for i in casename:
        casestr = casename2number(i)
        pictitle = f"{casestr}_l-{level}_p-{parent}"

        basename = os.path.basename(i).split("-")[0]
        data = pd.DataFrame(childdatadict[basename])
        data  = data.sort_index(axis=1)
        data  = data.sort_values(by='np', ascending=True)
        data.set_index(["np"], inplace=True)
        # print("basename",basename)

        # x,y,legend_name,Totaltime = getpltdata(data)

        if plttype == "time_heatmap" :
            plot_heatmap(savedir, pictitle + "-timeheatmap", f"{plttype}(s)", data.T)
        elif plttype == "rank_heatmap" :
            data = data.T
            for i in list(data.columns):
                data[i] = data[i].rank()
            plot_heatmap(savedir, pictitle + "-rankheatmap",plttype,data)
        # elif plttype == "time_linechart" :
            # plot_linechart(savedir, pictitle, "Time(s)" ,data)
            # if(level == 0):
            #     plot_linechart(savedir,x,Totaltime,f"{plttype}(s)",legend_name,title)
            # else:
            #     plot_linechart(savedir,x,y,f"{plttype}(s)",legend_name,title)
        elif plttype == "pic_paperline" :
            casenumber = int(casestr.split("C")[-1]) - 1
            mypaper_subplot(axes[casenumber//3, casenumber%3], pictitle ,data)
        elif plttype == "pic_paperpercent":
            casenumber = int(casestr.split("C")[-1]) - 1
            my_paper_plot_percent(axes[casenumber//3, casenumber%3], pictitle ,data)
        # elif plttype == "percent_par" or plttype == "percent_total":
        #     plot_percent(savedir,x,y,plttype,legend_name,title,Totaltime,1)
        # elif plttype == "efficiency" :
        #     x,y = getplteffdata(x,y)
        #     plot_linechart(savedir,x,y,f"{plttype}(%)",legend_name,title+"-efficiency")
        # elif plttype == "sizeperprocess" :
        #     plot_sizeperprocess(savedir,str(basename),x,title+"-sizeperprocess")
    if plttype == "pic_paperline" or plttype == "pic_paperpercent" :
        close_sublot(fig,axes,savedir, pictitle ,plttype)
# def test():
#     level = 3
#     parent = "CreateS"
#     file = "/home/mrz/data/result_23_1_13_tian/result-20230112T103924/jxpamg_solver01/jxpamg_solver01-database.json"
#     with open(file, "r") as fo:
#         text = fo.read()
#         jsondata = json.loads(text)
#     getavetime(jsondata, level, parent)

def main():
    parser = make_parser()
    args = parser.parse_args()
    runcases = get_runcases(args.LOGDIR)
    casename = getloadcase(runcases, args.CASENAME)
    datadict = getjsondata(casename, args.out_level, args.component_parent, args.out_type)
    
    # print(datadict)
    # for i in datadict[0]:
    #     print(pd.DataFrame(datadict[0][i]))

    out_title = createtitle(args.out_level, args.component_parent)

    #TODO:下面代码待整理，要和规划的目标一致，保证代码简洁，单个函数功能明确
    out_type_name = args.out_type.split("_")[0]
    if(out_type_name == "csv"):
        save2csv(args.savedir, casename, out_title, args.out_type, datadict)
    if(out_type_name == "pic"):
        save2pic(args.savedir, casename, args.out_level, args.component_parent, args.plot_type, datadict)


if __name__ == "__main__":
    main()
    # test()
